# -*- coding: utf-8 -*-
# File : service.py
# Author: taoyahui
# Date : 2021/6/16

import sys

sys.path.append('..')
from utils.path_util import get_root_dir

sys.path.append(get_root_dir())
from utils import mysql_util, common_util
from tasks.secet_free import key_distribute
from tasks.init_sys_env import get_init_sys_env, init_hostname, sync_remote_source
from tasks.patch_deploy import execute_task
from tasks import file_distribute
import time
import results
from application_factory.models import StrategyCls
from application_factory.extensions import scheduler
from application_factory.vul_scan_service import vul_scan_service, vul_scan_immediately
import json
from application_factory.models import InventoryCls
from utils import common_util, job_util
from datetime import datetime
import results
from utils.log_util import logger


# 资产初始化配置
# 需要配置内容：登录用户名和密码，
# 从通过sshpass免密配置
# 通过ansible远程登录获取用户实际操作系统
def linux_inventory_init_configure(inventory_ip, username, password, port=22):
    """
    资产初始化配置
    需要配置内容：登录用户名和密码，
    从通过sshpass免密配置
    通过ansible远程登录获取用户实际操作系统
    :param inventory_ip: 资产ip
    :param username: 用户名
    :param password: 密码
    :return:
    """
    conn = mysql_util.get_conn()
    cursor = mysql_util.get_cursor(conn)
    inventory_discover_sql = f"select * from inventory where ip = '{inventory_ip}'"
    result = mysql_util.exec_sql_new(inventory_discover_sql)
    if result is not None:
        inventory_cls = InventoryCls(result[0])
        # inventory = result[0]  # inventory[0] -> ip, inventory[1] -> os
        logger.info(f"linux_inventory_init_configure -> ip: {inventory_cls.ip} 开始初始化！!")

        ip = inventory_cls.ip
        os = inventory_cls.os
        hostname = inventory_cls.hostname

        # 如果hostname为空，表示没有
        if hostname is None:
            return results.get_hostname_error_result

        save_inventory_cert(username, password, ip)

        # 判断nmap扫到的资产是不是Linux操作系统
        if os == 'None' or os is None or 'Linux' in os:
            # 首先，免密登录
            result = key_distribute(ip, username, password)
            if result is not results.success_result:
                return result

            # 然后，获取操作系统信息并更新
            os = get_init_sys_env(ip, username, password, port)
            update_inventory_os(ip, os)
            # print(os)

            # # 设置hostname
            # hostname = init_hostname(ip, username, password, 22)
            # update_inventory_hostname(ip, hostname)
            logger.info(f"linux_inventory_init_configure -> ip: {inventory_cls.ip} 开始扫描!")
            vul_scan_immediately(inventory_cls.id)
            sync_remote_source_service(ip, username, password, port)
            # 开始定义执行漏扫的默认job(每天的02:00执行)

            vul_scan_service(ip, hostname, os)
    mysql_util.release(conn, cursor)
    return results.success_result


# 保存传输的用户名和密码
# 参数: username  -> 用户名, password -> 密码, ip ...
def save_inventory_cert(username, password, ip):
    save_sql = f'Update inventory set username = "{username}", password = "{password}" where ip = "{ip}"'
    mysql_util.exec_sql_new(save_sql)


# 更新资产的操作系统内容
# 参数: ip -> 目标IP， os -> 目标操作系统
def update_inventory_os(ip, os):
    os = os[:30]
    # 添加is_available = 1表示资产可用
    update_sql = f"Update inventory set os = '{os}', is_available = 1 where ip = '{ip}'"
    mysql_util.exec_sql_new(update_sql)


# 更新资产的主机名
# 参数:
#   ip: 资产ip
#   hostname: 主机名
def update_inventory_hostname(ip, hostname):
    update_sql = f"Update inventory set hostname = '{hostname}' where ip = '{ip}'"
    mysql_util.exec_sql_new(update_sql)


def modify_strategy_service(s_id):
    if s_id is None:
        return results.req_info_incomplete_result
    res = job_util.strategy_job_find(s_id)
    if 'code' in res:
        return res
    return add_strategy_service(s_id)


# 添加策略Service
# 参数:
#       s_id: 策略编号
def add_strategy_service(s_id):
    """
    添加策略Service
    :param s_id: 策略编号(id)
    :return:
    """
    if s_id is None:
        return results.strategy_not_exist_result
    # 通过查表，确认策略内容
    get_strategy_sql = f"select * from strategy_manage where id = '{s_id}'"
    strategy_list = mysql_util.exec_sql_new(get_strategy_sql)
    if len(strategy_list) == 0:
        return results.strategy_not_exist_result
    strategy = strategy_list[0]
    add_strategy = StrategyCls(strategy)

    # 测试补丁分发
    if add_strategy.patch_distribute == 1:
        # 选择测试补丁
        is_test = 1
        # type ==2 为循环策略
        if add_strategy.strategy_type == 2:
            return add_strategy_job(add_strategy, is_test, 'cron')
        else:
            return add_strategy_job(add_strategy, is_test, 'date')

    # 正式补丁分发
    if add_strategy.patch_distribute == 2:
        # 选择测试补丁
        is_test = 0
        # type ==2 为循环策略
        if add_strategy.strategy_type == 2:
            return add_strategy_job(add_strategy, is_test, 'cron')
        else:
            return add_strategy_job(add_strategy, is_test, 'date')
        pass

    # 人工补丁分发
    if add_strategy.patch_distribute == 3:
        return add_strategy_job(add_strategy, None, 'immediately')

    # 文件分发
    if add_strategy.patch_distribute ==4:
        return file_distribute.send_file_lib(add_strategy)


# 策略添加的任务(job)
# 参数:
#   add_strategy: 添加的策略(实体)
#   is_test: 是否策略对应的主机为测试机
#   trigger_type: 策略执行的类型. date --> 只执行一次   cron --> 执行若干次
def add_strategy_job(add_strategy, is_test=None, trigger_type='date'):
    """
    策略添加的任务(job)
    :param add_strategy: 添加的策略(实体)
    :param is_test: 是否策略对应的主机为测试机
    :param trigger_type: 策略执行的类型. date --> 只执行一次   cron --> 执行若干次
    :return:
    """
    task_name = f"Strategy-{add_strategy.strategy_id}-{add_strategy.strategy_name}"
    if trigger_type == 'date':
        exec_time_json = json.loads(add_strategy.exec_time)
        if 'mode' not in exec_time_json or exec_time_json['mode'] != 1:
            return results.strategy_exec_time_error_result
        time = exec_time_json['time']
        cron_rel = dict(run_date=datetime.strptime(time, '%Y-%m-%d %H:%M'))
        scheduler.add_job(func=execute_strategy,
                          id=task_name,
                          kwargs={'strategy_id': add_strategy.strategy_id, 'is_test': is_test},
                          trigger=trigger_type, **cron_rel,
                          replace_existing=True)
        # 记录执行的任务相关信息（时间，周期）
        job_util.add_strategy_job_record(task_name, cron_rel, 2, add_strategy.strategy_id)
    elif trigger_type == 'cron':
        exec_time_json = json.loads(add_strategy.exec_time)
        if 'mode' not in exec_time_json or exec_time_json['mode'] != 2:
            return results.strategy_exec_time_error_result
        week = int(exec_time_json['week']) - 1
        time_str = exec_time_json['time']
        tm_hour, tm_min = common_util.time_hour_min_format(time_str)
        cron_rel = dict(minute=f"{tm_min}", hour=f"{tm_hour}", day_of_week=week)
        scheduler.add_job(func=execute_strategy,
                          id=f"Strategy-{add_strategy.strategy_id}-{add_strategy.strategy_name}",
                          kwargs={'strategy_id': add_strategy.strategy_id, 'is_test': is_test},
                          trigger=trigger_type, **cron_rel,
                          replace_existing=True)
        # 记录执行的任务相关信息（时间，周期）
        job_util.add_strategy_job_record(task_name, cron_rel, 0, add_strategy.strategy_id)
    elif trigger_type == 'immediately':
        scheduler.add_job(func=execute_strategy,
                          id=f"Strategy-immediately-{add_strategy.strategy_id}-{add_strategy.strategy_name}",
                          kwargs={'strategy_id': add_strategy.strategy_id, 'is_test': is_test})
    return results.success_result


def get_exec_cron(add_strategy):
    """
    获取执行策略的执行模式--> 按天、
    :param add_strategy:
    :return:
    """
    exec_time_json = json.loads(add_strategy.exec_time)

    if 'mode' not in exec_time_json:
        return results.strategy_exec_time_incomplete_result
        # week = exec_time_json['week']

    # 一次性执行的策略
    if exec_time_json['mode'] == 1:
        time = exec_time_json['time']
        return time
    if exec_time_json['mode'] == 2:
        week = int(exec_time_json['week']) - 1
        time_str = exec_time_json['time']
        tm_hour, tm_min = common_util.time_hour_min_format(time_str)
        return week, tm_hour, tm_hour


# 执行策略
# 参数: strategy_id -> 策略编号
def execute_strategy(strategy_id, is_test=None):
    """
    执行策略
    :param strategy_id: 策略的id
    :param is_test: 是否执行资产组中的测试机器
    :return:
    """
    timestamp = int(time.time())
    get_strategy_sql = f"select * from strategy_manage where id = '{strategy_id}'"
    res = mysql_util.exec_sql_new(get_strategy_sql)
    if len(res) == 0:
        return results.strategy_not_exist_result
    strategy = StrategyCls(res[0])
    update_strategy_state(strategy, 1)

    # 如果分发类型为4，说明是文件分发
    if strategy.patch_distribute == 4:
        return file_distribute.send_file_lib(strategy)
    else:
        # 执行策略
        res = execute_task(strategy_id=strategy_id,
                     patch_group_id=strategy.patch_group_id,
                     inventory_group_id=strategy.inventory_group_id,
                     strategy_name=strategy.strategy_name,
                     timestamp=timestamp,
                     is_test=is_test)
        if res['code'] == 200:
            update_strategy_state(strategy, 2)


def update_strategy_state(strategy, state):
    """
    更新策略执行状态
    :param strategy:
    :param state: 0-> 未开始 1->执行中 2->已完成
    :return:
    """
    update_sql = f"UPDATE strategy_manage set execution_status = {state} where id = {strategy.strategy_id}"
    mysql_util.exec_sql_new(update_sql)


def show_strategy_service():
    jobs = scheduler.get_jobs()
    job_list = []
    for job in jobs:
        job_dict = {
            'name': job.name,
            'next_run_time': common_util.specify_time_format(job.next_run_time),
            'cron': str(job.trigger)
        }
        job_list.append(job_dict)

    return job_list


# 删除策略Service
def remove_strategy_service(s_id):
    """
    删除策略Service
    :param s_id:
    :return:
    """
    get_strategy_sql = f"select * from strategy_manage where id = '{s_id}'"
    strategy_list = mysql_util.exec_sql_new(get_strategy_sql)
    if len(strategy_list) == 0:
        return results.strategy_not_exist_result
    remove_strategy = StrategyCls(strategy_list[0])
    res = job_util.remove_strategy_execute_job(s_id)
    if res['code'] != 200:
        return res
    scheduler.remove_job(f"Strategy-{remove_strategy.strategy_id}-{remove_strategy.strategy_name}")
    return results.success_result


def sync_remote_source_service(ip, username, password, port):
    logger.info(f"开始设置远程源:IP -> {ip}")
    job = scheduler.add_job(func=sync_remote_source,
                            id=f"Sync-Remote_source-{ip}",
                            kwargs={'ip': ip, 'username': username, 'passwd':password, 'port':port}
                            )


if __name__ == '__main__':
    # linux_inventory_init_configure('192.168.50.251', 'root', '123456')
    execute_strategy(48)
    # add_strategy_service(14)
